import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import spsolve
import os


def BarsLength(x, y, z, ele):#（计算数组长度）
    return np.sqrt((x[ele[:, 1] - 1] - x[ele[:, 0] - 1]) ** 2 +
                   (y[ele[:, 1] - 1] - y[ele[:, 0] - 1]) ** 2 +
                   (z[ele[:, 1] - 1] - z[ele[:, 0] - 1]) ** 2)

def CoordTransform(x, y, z, L):#(计算坐标变换矩阵
    x1, x2 = x
    y1, y2 = y
    z1, z2 = z
    if x1 == x2 and y1 == y2:#（布尔“与”，都为真则真）
        if z2 > z1:
            Lambda = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
        else:
            Lambda = np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]])
    else:
        CXx = (x2 - x1) / L#（三个方向余弦，局部坐标系相对于全局坐标系方向）
        CYx = (y2 - y1) / L
        CZx = (z2 - z1) / L
        D = np.sqrt(CXx * CXx + CYx * CYx)#（点1和点2形成的向量在xy平面上的投影长度）
        CXy = -CYx / D#计算局部y轴的单位向量
        CYy = CXx / D
        CZy = 0
        CXz = -CXx * CZx / D#（计算局部z轴的单位向量）
        CYz = -CYx * CZx / D
        CZz = D
        Lambda = np.array([[CXx, CYx, CZx],#（旋转矩阵的列向量[CXx, CYx, CZx]，分别代表局部坐标系的基向量在全局坐标系中的表示）
                           [CXy, CYy, CZy],
                           [CXz, CYz, CZz]])
    R = np.block([#(R 为12x12矩阵)
        [Lambda, np.zeros((3, 9))],#（旋转矩阵和3x9的零矩阵）
        [np.zeros((3, 3)), Lambda, np.zeros((3, 6))],#（3x3的零矩阵、旋转矩阵和3x6的零矩阵）
        [np.zeros((3, 6)), Lambda, np.zeros((3, 3))],
        [np.zeros((3, 9)), Lambda]
    ])
    return R

def Bernoulli3DElementKe(prop, R, L, icoord):
    E, A, Iy, Iz, G, J = prop
    k1 = E * A / L#(轴向刚度系数）
    k2 = 12 * E * Iz / (L ** 3)
    k3 = 6 * E * Iz / (L ** 2)
    k4 = 4 * E * Iz / L
    k5 = 2 * E * Iz / L
    k6 = 12 * E * Iy / (L ** 3)
    k7 = 6 * E * Iy / (L ** 2)
    k8 = 4 * E * Iy / L
    k9 = 2 * E * Iy / L
    k10 = G * J / L
    ke = np.array([
        [k1, 0, 0, 0, 0, 0, -k1, 0, 0, 0, 0, 0],
        [0, k2, 0, 0, 0, k3, 0, -k2, 0, 0, 0, k3],
        [0, 0, k6, 0, -k7, 0, 0, 0, -k6, 0, -k7, 0],
        [0, 0, 0, k10, 0, 0, 0, 0, 0, -k10, 0, 0],
        [0, 0, -k7, 0, k8, 0, 0, 0, k7, 0, k9, 0],
        [0, k3, 0, 0, 0, k4, 0, -k3, 0, 0, 0, k5],
        [-k1, 0, 0, 0, 0, 0, k1, 0, 0, 0, 0, 0],
        [0, -k2, 0, 0, 0, -k3, 0, k2, 0, 0, 0, -k3],
        [0, 0, -k6, 0, k7, 0, 0, 0, k6, 0, k7, 0],
        [0, 0, 0, -k10, 0, 0, 0, 0, 0, k10, 0, 0],
        [0, 0, -k7, 0, k9, 0, 0, 0, k7, 0, k8, 0],
        [0, k3, 0, 0, 0, k5, 0, -k3, 0, 0, 0, k4]
    ])

    if icoord == 1:
        Ke = R.T @ ke @ R
    elif icoord == 2:
        Ke = ke

    return Ke

def Bernoulli3DFEA(Node, Element, Material, force, constrain):
    Dof = 6#（每个节点自由度为6，三个平移，三个 旋转）
    EDof = Dof * 2#（每个单元两个节点，12个自由度）
    Dofs = Dof * Node.shape[0]#（计算整个结构总自由度数，单个节点的自由度数乘节点数组中节点个数）
    EleCount = Element.shape[0]#（获取单元数量）
    KKG = sp.lil_matrix((Dofs, Dofs))#（矩阵储存结构的全局刚度矩阵）
    FFG = sp.lil_matrix((Dofs, 1))#（矩阵存储整个结构的载荷向量，列向量长度等于总自由度数）
    U = sp.lil_matrix((Dofs, 1))#（矩阵存储求解后的位移向量）
    x = Node[:, 0]#（从节点数组中提取所有节点的x坐标）
    y = Node[:, 1]#（从节点数组中提取所有节点的y坐标）
    z = Node[:, 2]#（从节点数组中提取所有节点的z坐标）
    BarLength = BarsLength(x, y, z, Element)#（计算数组长度）

    for iEle in range(EleCount):#（遍历每个单元）
        n1, n2 = Element[iEle] - 1#（提取当前单元的前两个节点编号并-1，因为索引从0开始）
        R = CoordTransform([x[n1], x[n2]], [y[n1], y[n2]], [z[n1], z[n2]], BarLength[iEle])
        ElementStiffnessMatrix = Bernoulli3DElementKe(Material[0], R, BarLength[iEle], 1)
        n1_dofs = np.arange(n1 * Dof, (n1 + 1) * Dof)#（定义数组第一个节点自由度索引，从第一个节点的自由度开始，到第二个节点的自由度截止，中间包含的就是第一个节点的所有自由度索引）
        n2_dofs = np.arange(n2 * Dof, (n2 + 1) * Dof)#（定义数组第二个节点自由度索引，从第一个节点的自由度开始，到第二个节点的自由度截止，中间包含的就是第二个节点的所有自由度索引）
        ElementNodeDOF = np.concatenate((n1_dofs, n2_dofs))#（将两个自由度节点合并成一个数组,一个单元的节点自由度索引）


        for i in range(len(ElementNodeDOF)):
            for j in range(len(ElementNodeDOF)):
                KKG[ElementNodeDOF[i], ElementNodeDOF[j]] += ElementStiffnessMatrix[i, j]#（将元素的局部刚度矩阵组装到全局刚度矩阵 KKG 中）

    if force.size:#（检查数组force是否为零，若是空数组则条件为假，不会执行）
        for i in range(len(force)):#（遍历循环force数组元素，生成一个从0到force数组长度-1的序列）
            m = int(force[i, 0])#（从force数组提取第一个值并转为整数，代表受力节点编号）
            n = int(force[i, 1])#（从force数组提取第二个值并转为整数，代表节点m上力的自由度编号）
            value = force[i, 2]#（从force数组提取第三个值，代表节点m的自由度n上力的大小）
            FFG[Dof * (m - 1) + n - 1, 0] = value#（计算全局力向量FFG对应与节点m和自由度n的位置）

    for i in range(len(constrain)):#处理约束条件
        n = int(constrain[i, 0])  # 节点号
        d = int(constrain[i, 1])  # 局部自由度号
        m = (n - 1) * Dof + d - 1  # 全局自由度索引位置
        FFG = FFG - constrain[i, 2] * KKG[:, m]  # 载荷列阵均减去相应元素
        FFG[m, 0] = constrain[i, 2]  # 载荷列阵相应位置更改为已知位移
        KKG[:, m] = 0  # 刚度矩阵列置零，取m+1列为零
        KKG[m, :] = 0  # 刚度矩阵行置零，取m+1行为零
        KKG[m, m] = 1.0  # 刚度矩对角元素阵置 1

    U = spsolve(KKG.tocsc(), FFG)

    return U
